package linker

import (
	"debug/elf"
	"math"
	"risc-v/ld/pkg/utils"
	"sort"
)

const NeedGotYp uint32 = 1 << 0

// ResolveSymbols 解析 ObjectFile 中的符号
func ResolveSymbols(ctx *Context) {
	// 解析所有的 ObjectFile 中的符号
	for _, file := range ctx.ObjectFiles {
		file.ResolveSymbol()
	}

	// 标记活跃的 ObjectFile
	MarkLiveObjects(ctx)

	// 清除非活跃的 ObjectFile 中的的 Global 符号
	for _, file := range ctx.ObjectFiles {
		if !file.IsAlive {
			file.ClearSymbols()
		}
	}

	// 移除非活跃的 ObjectFile
	ctx.ObjectFiles = utils.RemoveIf(ctx.ObjectFiles, func(file *ObjectFile) bool {
		return !file.IsAlive
	})

}

// MarkLiveObjects 标记活跃的 ObjectFile
func MarkLiveObjects(ctx *Context) {
	roots := make([]*ObjectFile, 0)
	//读入最初的 Object 文件 也就是最初的命令行上给出的 ObjectFile
	for _, file := range ctx.ObjectFiles {
		if file.IsAlive {
			roots = append(roots, file)
		}
	}

	utils.Assert(len(roots) > 0)

	for len(roots) > 0 {
		file := roots[0]

		if !file.IsAlive {
			continue
		}

		// 递归的去标记所有的活跃的 ObjectFile
		file.MarkLiveObjects(func(file *ObjectFile) {
			roots = append(roots, file)
		})

		roots = roots[1:]
	}
}

func RegisterSectionPieces(ctx *Context) {
	for _, file := range ctx.ObjectFiles {
		file.RegisterSectionPieces()
	}
}

// CreateSyntheticSections 创建合成节 如输出的 ELF 文件头 程序头等
func CreateSyntheticSections(ctx *Context) {
	push := func(chunk Chunker) Chunker {
		ctx.Chunks = append(ctx.Chunks, chunk)
		return chunk
	}

	// 创建 ELF 文件头
	ctx.ELFHeader = push(NewOutPutELFHeader()).(*OutPutELFHeader)
	// 创建程序头
	ctx.ProgramHeader = push(NewOutPutProgramHeader()).(*OutPutProgramHeader)
	// 创建Section Header
	ctx.SectionHeader = push(NewOutPutSectionHeader()).(*OutPutSectionHeader)
	// 创建全局偏移量表
	ctx.Got = push(NewGotSections()).(*GotSection)
}

func SetOutPutSectionOffset(ctx *Context) uint64 {
	addr := IMAGE_BASE // 虚拟地址的开始

	// 遍历整个Chunks
	for _, chunk := range ctx.Chunks {
		// 不需要分配地址
		if chunk.GetSectionHeader().Flags&uint64(elf.SHF_ALLOC) == 0 {
			continue
		}

		// 地址对齐
		addr = utils.AlignTo(addr, chunk.GetSectionHeader().AddrAlign)
		chunk.GetSectionHeader().Addr = addr

		// .tbss 不占用磁盘空间,这里要跳过
		if !isTbss(chunk) {
			// 地址累加
			addr += chunk.GetSectionHeader().Size
		}
	}

	i := 0
	first := ctx.Chunks[0]

	for {
		sectionHeader := ctx.Chunks[i].GetSectionHeader()
		sectionHeader.Offset = sectionHeader.Addr - first.GetSectionHeader().Addr

		i++

		if i >= len(ctx.Chunks) ||
			ctx.Chunks[i].GetSectionHeader().Flags&uint64(elf.SHF_ALLOC) == 0 {
			break
		}
	}

	lastShdr := ctx.Chunks[i-1].GetSectionHeader()
	fileOff := lastShdr.Offset + lastShdr.Size

	for ; i < len(ctx.Chunks); i++ {
		sher := ctx.Chunks[i].GetSectionHeader()
		fileOff = utils.AlignTo(fileOff, sher.AddrAlign)
		sher.Offset = fileOff // offset 初始化
		fileOff += sher.Size
	}

	ctx.ProgramHeader.UpdateSectionHeader(ctx)
	return fileOff

}

// BindSection 填充 OutPutSection 中的 Members(InPutSection) 成员
// 遍历 Context 中每个 ObjectFile 的 Section 将其和 OutPutSection 关联起来
func BindSection(ctx *Context) {
	group := make([][]*InputSection, len(ctx.OutPutSections))
	for _, file := range ctx.ObjectFiles {
		for _, insection := range file.Sections {
			if insection == nil || !insection.IsAlive {
				continue
			}
			// 获取 OutPutSection 中的索引
			index := insection.OutPutSection.Index
			// 加入对应 OutPutSection 中的 group 中
			group[index] = append(group[index], insection)
		}
	}

	for index, outSection := range ctx.OutPutSections {
		outSection.Members = group[index]
	}
}

// CollectOutPutSections 收集 context 中所有有效的 OutPutSection
func CollectOutPutSections(ctx *Context) []Chunker {
	outSections := make([]Chunker, 0)
	for _, section := range ctx.OutPutSections {
		if len(section.Members) > 0 {
			outSections = append(outSections, section)
		}
	}

	// 收集合并节
	for _, osec := range ctx.MergedSections {
		if osec.SectionHeader.Size > 0 {
			outSections = append(outSections, osec)
		}
	}

	return outSections
}

// ComputeSectionSizes 计算 OutPutSection 中的 Size 和 地址对齐量
// 这里进行的类似于把来自多个 object 文件中的 .text 段的数据合并到一起
func ComputeSectionSizes(ctx *Context) {
	// 处理每一个 OutPutSection
	for _, osec := range ctx.OutPutSections {
		offset := uint64(0)
		p2align := int64(0)
		// 遍历一个 OutPutSection 中的所有的 InPutSection
		for _, isec := range osec.Members {
			offset = utils.AlignTo(offset, 1<<isec.P2Align)                    //计算对齐后的offset
			isec.Offset = uint32(offset)                                       // 赋值对齐后的 offset
			offset += uint64(isec.SectionSize)                                 //更新 offset
			p2align = int64(math.Max(float64(p2align), float64(isec.P2Align))) //更新对齐
		}
		osec.SectionHeader.Size = offset
		// OutPutSection 中的 AddrAlign 取所有的 InPutSection 中的最大值
		osec.SectionHeader.AddrAlign = 1 << p2align // p2align 是 2 的幂数，直接左移即可
	}
}

func SortOutPutSections(ctx *Context) {

	// 标记 rank 值
	// rank 值越小排序就越前面
	rank := func(chunk Chunker) int32 {
		typ := chunk.GetSectionHeader().Type
		flags := chunk.GetSectionHeader().Flags

		if flags&uint64(elf.SHF_ALLOC) == 0 {
			return math.MaxInt32 - 1
		}
		if chunk == ctx.SectionHeader {
			return math.MaxInt32
		}

		if chunk == ctx.ELFHeader {
			return 0
		}
		if chunk == ctx.ProgramHeader {
			return 1
		}

		if typ == uint32(elf.SHT_NOTE) {
			return 2
		}

		bool2int := func(b bool) int {
			if b {
				return 1
			}
			return 0
		}

		writeable := bool2int(flags&uint64(elf.SHF_WRITE) != 0)
		notExec := bool2int(flags&uint64(elf.SHF_EXECINSTR) == 0)
		notTls := bool2int(flags&uint64(elf.SHF_TLS) == 0)
		// TLS  thread local storage
		// 线程自身存储的数据
		isBss := bool2int(typ == uint32(elf.SHT_NOBITS))

		return int32(writeable<<7 | notExec<<6 | notTls<<5 | isBss<<4)
	}

	sort.SliceStable(ctx.Chunks, func(i, j int) bool {
		return rank(ctx.Chunks[i]) < rank(ctx.Chunks[j])
	})
}
func ScanRelocations(ctx *Context) {
	for _, file := range ctx.ObjectFiles {
		file.ScanRelocations()
	}

	symbols := make([]*Symbol, 0)
	for _, file := range ctx.ObjectFiles {
		for _, symbol := range file.Symbols {
			if symbol.File == file && symbol.Flags != 0 {
				symbols = append(symbols, symbol)
			}
		}
	}

	for _, symbol := range symbols {
		if symbol.Flags&NeedGotYp != 0 {
			ctx.Got.AddGotTpSymbol(symbol)
		}

		symbol.Flags = 0
	}
}

func ComputeMergedSectionSize(ctx *Context) {
	for _, osec := range ctx.MergedSections {
		osec.AssignOffsets()
	}
}

// 判断是否是 .tbss section
// .data .bss
// .tdata .tbss
func isTbss(chunk Chunker) bool {
	sectionHeader := chunk.GetSectionHeader()
	return sectionHeader.Type == uint32(elf.SHT_NOBITS) &&
		sectionHeader.Flags&uint64(elf.SHF_TLS) != 0
}
